package com.lun.swordtowardoffer2.c09;

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Queue;

public class KSmallestPairs {

	//方法一：使用最大堆
	public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k){
		Queue<int[]> maxHeap = new PriorityQueue<>((p1, p2)
				-> p2[0] + p2[1] - p1[0] - p1[1]);
		
		for(int i = 0; i < Math.min(k, nums1.length); i++) {
			for(int j = 0; j < Math.min(k, nums2.length); j++) {
				
				if(maxHeap.size() < k) {
					maxHeap.offer(new int[] {nums1[i], nums2[j]});
				}else {
					int[] root = maxHeap.peek();
					if(nums1[i] + nums2[j] < root[0] + root[1]) {
						maxHeap.poll();
						maxHeap.offer(new int[] {nums1[i], nums2[j]});
					}
				}				
			}
		}
		
		List<int[]> result = new ArrayList<>();
		for(int[] pair : maxHeap) {
			result.add(pair);
		}
		return result.stream().sorted((p1, p2)
				-> p1[0] + p1[1] - p2[0] - p2[1]).toList();
	}
	
	//方法二：使用（候选下标）最小堆
	public List<int[]> kSmallestPairs2(int[] nums1, int[] nums2, int k){
		Queue<int[]> candidateIndexPairMinHeap = new PriorityQueue<>((indexPair1, indexPair2)//
				-> nums1[indexPair1[0]] + nums2[indexPair1[1]] - //
						nums1[indexPair2[0]] - nums2[indexPair2[1]]);
		
		if(nums2.length > 0) {
			for(int i = 0; i < Math.min(k, nums1.length); i++) {
				candidateIndexPairMinHeap.offer(new int[] {i, 0});
			}
		}
		
		List<int[]> result = new ArrayList<>();
		while(k-- > 0 && !candidateIndexPairMinHeap.isEmpty()) {
			int[] indexPair = candidateIndexPairMinHeap.poll();
			result.add(new int[] {nums1[indexPair[0]], nums2[indexPair[1]]});
			
			if(indexPair[1] < nums2.length) {
				candidateIndexPairMinHeap.offer(new int[] {indexPair[0], indexPair[1] + 1});
			}
		}
		return result;
	}

}
